#!/usr/bin/env bash

###############################################################
# File Name: install_es_cluster.sh
# Version: V1.0
# Author: Brian Hsiung
# Created Time : 2025-01-21 9:37:53
# Description: elasticsearch 7.16.3 三节点集群部署脚本
##############################################################

# 调试跟踪（-e: 遇错退出， -u: 捕获未定义的变量， -x: 调试模式，在执行每个命令前，将其打印到标准错误输出）
# set -eux (不可使用-e) 
# set -x

# ES 节点名(留空,下文自动赋值)
ES_NODE_NAME=
# ES 集群状态(留空,下文自动赋值)
ES_HEALTH_STATUS=
# 对外服务端口
ES_HTTP_PORT=9200
# 集群通讯端口
ES_CLUSTER_PORT=9300
# ES版本
ES_VERSION="7.16.3"
# ES用户
ES_USER="elkuser"
# ES 运行内存
ES_JAVA_OPTS="4g"
# ES 安装目录
ES_BASE_DIR="/data/elasticsearch"
# ES 日志目录
ES_LOG="$ES_BASE_DIR/logs"
# ES 数据目录
ES_DATA="$ES_BASE_DIR/data"
# ES 集群名称
ES_CLUSTER_NAME="elasticsearch-cluster"
# 软连接目录，若ES的BASE目录为/data/elasticsearch/elasticsearch-7.16.3，ES_HOME为/data/elasticsearch/elasticsearch
ES_HOME="$ES_BASE_DIR/elasticsearch"
# ES 每个节点的名称
ES_NODE_NAME_LIST=(elastic01 elastic02 elastic03)
# 阿里云OSS存储地址
OSS_BASE_URL="https://brianhsiung.oss-cn-hangzhou.aliyuncs.com"

# ES 集群的IP列表, 此处需要根据实际情况修改
ES_IP_LIST=(172.26.85.157 172.26.85.158 172.26.85.159)

# # 使用关联数组定义节点拓扑
# declare -A ES_CLUSTER_NODES=(
#     ["172.26.85.157"]="elastic01"
#     ["172.26.85.158"]="elastic02"
#     ["172.26.85.159"]="elastic03"
# )

DIR="$(cd "$(dirname "$0")" && pwd)"
cd "$DIR" || exit 1

##############################################
# 日志函数
##############################################
log_info()    { echo -e "\033[34m[INFO] $*\033[0m"; }
log_warn()    { echo -e "\033[33m[WARNING] $*\033[0m"; }
log_success() { echo -e "\033[32m[SUCCESS] $*\033[0m"; }
log_error()   { echo -e "\033[31m[ERROR] $*\033[0m"; exit 1; }

##############################################
# 检查依赖
##############################################
check_dependency() {

    local -a required_cmds=("wget" "unzip" "tar" "groupadd" "useradd" "chkconfig")
    local missing=()

    log_info "检查依赖命令: ${required_cmds[*]}"
    for cmd in "${required_cmds[@]}"; do
        if ! command -v "$cmd" &>/dev/null; then
            missing+=("$cmd")
        fi
    done

    ((${#missing[@]})) && log_error "缺失依赖命令: ${missing[*]}"
}

##############################################
# 检查端口
##############################################
check_ports() {
    local port=$1
    if ss -lntu | grep -q ":$port "; then
        log_error "端口$port 已被占用"
    fi
}

##############################################
# 设置目录权限（递归检查上级目录）
##############################################
set_parent_permissions() {
    ### 某些等保三服务器默认新建的的目录权限为700，导致服务无法启动，所以需要遍历上级目录，将权限修改为755
    # 获取真实目录
    local target_dir
    target_dir=$(realpath "$1")
    # 获取上级目录
    local current_dir
    current_dir=$(dirname "$target_dir")

    # 做个判断，避免进入死循环
    [ -d "$target_dir" ] || log_error "目录不存在 [$target_dir]"

    # 排除敏感目录
    local -a exclude_dirs
    exclude_dirs=("/usr" "/usr/local" "/bin" "/sbin" "/lib" "/etc" "/dev" "/var" "/root" "/usr/share")

    while [[ "$current_dir" != "/" ]]; do

        # 排除敏感目录
        if printf '%s\n' "${exclude_dirs[@]}" | grep -qx "$current_dir"; then
            log_info "跳过系统保护目录: $current_dir"
            current_dir=$(dirname "$current_dir")
            continue
        fi

        # 跳过不存在的目录
        if [[ ! -d "$current_dir" ]]; then
            log_warn "警告: 目录不存在 [$current_dir]"
            current_dir=$(dirname "$current_dir")
            continue
        fi
        
        # 获取当前目录权限
        local perm_oct
        perm_oct=$(stat -c '%a' "$current_dir")
        log_info "目录 $current_dir 的原始权限为: $perm_oct"
        
        if [[ "$perm_oct" != "755" ]]; then
            log_info "正在设置权限: [755] -> $current_dir"
            if ! chmod 755 "$current_dir"; then
                log_error "错误: 无法设置目录权限 [$current_dir]"
            fi
        else
            log_success "权限正确 [755]: $current_dir"
        fi

        # 向上一级目录
        current_dir=$(dirname "$current_dir")
    done
}

##############################################
# 初始化环境
##############################################
init_env() {

    [[ $EUID -ne 0 ]] && log_error "请使用root账号执行此脚本"

    # 检查依赖
    check_dependency
    # 检查端口
    check_ports "$ES_HTTP_PORT"
    check_ports "$ES_CLUSTER_PORT"

    # 当前主机IP
    HOST_IP=$(hostname -I | awk 'NR==1 {print $1}')
    if echo "${ES_IP_LIST[@]}" | grep -w "$HOST_IP" &> /dev/null; then
        log_info "当前主机IP为: $HOST_IP"
    else
        log_error "获取的主机IP: $HOST_IP 不在 ${ES_IP_LIST[*]} 中, 请手动设置 HOST_IP"
    fi
    # 设置ES节点名称
    for (( i=0; i < "${#ES_IP_LIST[@]}"; i++ )); do
        if [ "$HOST_IP" = "${ES_IP_LIST[i]}" ]; then
            ES_NODE_NAME="${ES_NODE_NAME_LIST[i]}"
            log_info "ES当前节点名称: $ES_NODE_NAME"
            break
        fi
    done

    [[ -z "$ES_NODE_NAME" ]] && log_error "ES_NODE_NAME的值为空"

    # 创建用户
    if ! getent group "$ES_USER" &>/dev/null; then
        log_info "创建组: $ES_USER"
        groupadd "$ES_USER"
    fi
    if ! getent passwd "$ES_USER" &>/dev/null; then
        log_info "创建用户: $ES_USER"
        useradd -g "$ES_USER" "$ES_USER"
    fi
    # 禁用交换分区swap
    if [ -n "$(swapon --show)" ]; then
        swapoff -a
        sed -i 's/.*swap.*/#&/' /etc/fstab
    fi
    # 禁用transparent_hugepage
    echo never > /sys/kernel/mm/transparent_hugepage/enabled
    echo never > /sys/kernel/mm/transparent_hugepage/defrag
    if ! grep transparent_hugepage /etc/rc.d/rc.local >/dev/null 2>&1; then
        cat >> /etc/rc.d/rc.local << EOF
echo never > /sys/kernel/mm/transparent_hugepage/enabled
echo never > /sys/kernel/mm/transparent_hugepage/defrag
EOF
    fi
    chmod +x /etc/rc.d/rc.local

    # vm.max_map_count
    if ! grep -w 'vm.max_map_count' /etc/sysctl.conf &> /dev/null; then
        log_info "设置vm.max_map_count"
        sysctl -w vm.max_map_count=1048576
        echo "vm.max_map_count = 1048576" >> /etc/sysctl.conf
        sysctl -p >/dev/null
    fi

    # 文件描述符限制
    local limits_conf="/etc/security/limits.conf"
    if ! grep -q "${ES_USER}.*nofile" "${limits_conf}"; then
        log_info "设置文件描述符限制"
        cat << EOF >> "${limits_conf}"
${ES_USER} soft nofile 65535
${ES_USER} hard nofile 65535
${ES_USER} soft memlock unlimited
${ES_USER} hard memlock unlimited
EOF
    fi

    if systemctl is-active --quiet firewalld; then
        log_info "开放防火墙端口$ES_HTTP_PORT, $ES_CLUSTER_PORT"
        firewall-cmd --zone=public --add-port=$ES_HTTP_PORT/tcp --permanent
        firewall-cmd --zone=public --add-port=$ES_CLUSTER_PORT/tcp --permanent
        firewall-cmd --reload
    fi
}

##############################################
# 安装Elasticsearch
##############################################
install_elasticsearch() {

    # 创建目录
    if [ ! -d "$ES_DATA" ]; then
        mkdir -p "$ES_DATA"
    elif [ "$(ls -A "$ES_DATA")" ]; then
        log_error "$ES_DATA 目录不为空, 请手动清理"
    fi

    if [ ! -d "$ES_LOG" ]; then
        mkdir -p "$ES_LOG"
    elif [ "$(ls -A "$ES_LOG")" ]; then
        log_error "$ES_LOG 目录不为空, 请手动清理"
    fi

    # 设置上级目录权限
    set_parent_permissions "$ES_BASE_DIR"

    local arch
    local es_pkg
    local arch_type
    arch="$(uname -m)"

    case "$arch" in
        aarch64|arm64)
            arch_type="aarch64"
            es_pkg="elasticsearch-$ES_VERSION-linux-aarch64.tar.gz"
        ;;
        amd64|x86_64)
            arch_type="x64"
            es_pkg="elasticsearch-$ES_VERSION-linux-x86_64.tar.gz"
        ;;
        *)
            echo "不支持的CPU架构: $arch"
            exit 1
        ;;
    esac

    if curl -sIfo /dev/null --max-time 5 --connect-timeout 10 "$OSS_BASE_URL"/health; then
        access_mode="external"
    fi

    if [ ! -f "$es_pkg" ] && [ "$access_mode" = "external" ]; then
        log_info "通过阿里云OSS下载elasticsearch安装包"
        wget -q "$OSS_BASE_URL/deploy/elastic/$arch_type/$es_pkg" || log_error "下载elasticsearch安装包失败"
    fi

    local ik_pkg="elasticsearch-analysis-ik-$ES_VERSION.zip"
    local pinyin_pkg="elasticsearch-analysis-pinyin-$ES_VERSION.zip"

    if [ ! -f "$ik_pkg" ] && [ "$access_mode" = "external" ]; then
        log_info "通过阿里云OSS下载elasticsearch-analysis-ik安装包"
        wget -q "$OSS_BASE_URL/deploy/elastic/$arch_type/$ik_pkg" || log_error "下载 $ik_pkg 安装包失败"
    fi

    if [ ! -f "$pinyin_pkg" ] && [ "$access_mode" = "external" ]; then
        log_info "通过阿里云OSS下载elasticsearch-analysis-pinyin安装包"
        wget -q "$OSS_BASE_URL/deploy/elastic/$arch_type/$pinyin_pkg" || log_error "下载 $pinyin_pkg 安装包失败"
    fi

    [ -f "$es_pkg" ] || log_error "$es_pkg 安装包不存在"
    [ -f "$ik_pkg" ] || log_error "$ik_pkg 安装包不存在"
    [ -f "$pinyin_pkg" ] || log_error "$pinyin_pkg 安装包不存在"

    # 解压
    log_info "解压$es_pkg"
    tar -zxf "$es_pkg" -C "$ES_BASE_DIR" --no-same-owner || log_error "解压elasticsearch安装包失败"
    log_info "解压$ik_pkg"
    unzip -qo "$ik_pkg" -d "$ES_HOME"-"$ES_VERSION"/plugins/analysis-ik || log_error "解压elasticsearch-analysis-ik安装包失败"
    log_info "解压$pinyin_pkg"
    unzip -qo "$pinyin_pkg" -d "$ES_HOME"-"$ES_VERSION"/plugins/analysis-pinyin || log_error "解压elasticsearch-analysis-pinyin安装包失败"

    # 创建软连接
    cd "$ES_BASE_DIR" || exit 1
    if [ -L "elasticsearch" ]; then
        log_warn "已经存在软链接 elasticsearch, 更新软连接"
        ln -sfn elasticsearch-"$ES_VERSION" elasticsearch || log_error "软链接更新失败"
    elif [ -e "elasticsearch" ]; then
        log_error "当前目录下存在名为 elasticsearch 的文件或目录"
    else
        log_info "创建软链接 elasticsearch"
        ln -s elasticsearch-"$ES_VERSION" elasticsearch
    fi
    cd "$DIR" || exit 1

    # 设置ES_JAVA_OPTS
    # sed -i -e "s/-Xms1g/-Xms$ES_JAVA_OPTS/g" -e "s/-Xmx1g/-Xmx$ES_JAVA_OPTS/g" $ES_HOME/config/jvm.options
    log_info "设置ES_JAVA_OPTS: $ES_JAVA_OPTS"
    {
        echo "-Xms$ES_JAVA_OPTS"
        echo "-Xmx$ES_JAVA_OPTS"
    } > "$ES_HOME"/config/jvm.options.d/jvm.options

    # 处理discovery.seed_hosts和cluster.initial_master_nodes
    local discovery_seed_hosts=""
    local cluster_initial_master_nodes=""
    for ((i=0; i<${#ES_IP_LIST[@]}; i++)); do
        discovery_seed_hosts+="\"${ES_IP_LIST[$i]}:$ES_CLUSTER_PORT\","
    done
    for ((i=0; i<${#ES_NODE_NAME_LIST[@]}; i++)); do
        cluster_initial_master_nodes+="\"${ES_NODE_NAME_LIST[$i]}\","
    done
    # 去掉最后一个逗号
    discovery_seed_hosts=${discovery_seed_hosts%,}
    cluster_initial_master_nodes=${cluster_initial_master_nodes%,}

    #discovery.seed_hosts: ["${ES_IP_LIST[0]}:$ES_CLUSTER_PORT", "${ES_IP_LIST[1]}:$ES_CLUSTER_PORT", "${ES_IP_LIST[2]}:$ES_CLUSTER_PORT"]
    #cluster.initial_master_nodes: ["${ES_NODE_NAME_LIST[0]}", "${ES_NODE_NAME_LIST[1]}", "${ES_NODE_NAME_LIST[2]}"]

    log_info "设置配置文件: $ES_HOME/config/elasticsearch.yml"
    cat > "$ES_HOME"/config/elasticsearch.yml << EOF
cluster.name: $ES_CLUSTER_NAME
node.name: $ES_NODE_NAME
node.master: true
node.data: true
path.data: $ES_DATA
path.logs: $ES_LOG
bootstrap.memory_lock: true
network.host: 0.0.0.0
http.port: $ES_HTTP_PORT
transport.tcp.port: $ES_CLUSTER_PORT
discovery.seed_hosts: [$discovery_seed_hosts]
cluster.initial_master_nodes: [$cluster_initial_master_nodes]
gateway.recover_after_nodes: 2
gateway.expected_nodes: 3
gateway.recover_after_time: 5m
discovery.zen.fd.ping_timeout: 300s
discovery.zen.fd.ping_retries: 8
discovery.zen.fd.ping_interval: 30s
discovery.zen.ping_timeout: 180s
discovery.zen.minimum_master_nodes: 2
xpack.security.enabled: true
xpack.security.transport.ssl.enabled: true
xpack.security.transport.ssl.verification_mode: certificate
xpack.security.transport.ssl.keystore.path: elastic-certificates.p12
xpack.security.transport.ssl.truststore.path: elastic-certificates.p12
EOF

    chown "$ES_USER":"$ES_USER" -R "$ES_BASE_DIR"
    find "$ES_HOME"/bin -type f -exec chmod 755 {} \;
    find "$ES_HOME"/jdk/bin -type f -exec chmod 755 {} \;
    # # 使用普通用户启动ES
    # su - $ES_USER -c "$ES_HOME/bin/elasticsearch -d"
}

add_service_config() {
    log_info "创建启动脚本 /etc/rc.d/init.d/elasticsearch"
    cat > /etc/rc.d/init.d/elasticsearch << EOF
#!/bin/bash
#chkconfig:2345 30 70
#description:elasticsearch

ES_HOME="$ES_HOME"

case \$1 in
    start)
        su - $ES_USER -c "\$ES_HOME/bin/elasticsearch -d"
        echo "elasticsearch is started"
        ;;
    stop)
        pgrep -f "\$ES_HOME" | xargs kill -9
        echo "elasticsearch is stopped"
        ;;
    restart)
        \$0 stop
        \$0 start
        echo "elasticsearch is started"
        ;;
    *)
        echo "start|stop|restart" 
        ;;
esac
EOF
    chmod +x /etc/rc.d/init.d/elasticsearch
    chkconfig --add elasticsearch
    chkconfig elasticsearch on
}

##############################################
# 等待集群启动
##############################################
wait_es_start() {
    local start_info
    local node_count=${#ES_IP_LIST[@]}
    local healthy_node_count=0

    while :
    do
        # 初始化健康节点计数器
        healthy_node_count=0

        for ip in "${ES_IP_LIST[@]}"; do
            log_info "检查节点 $ip 的状态..."
            # 获取节点健康信息
            start_info=$(curl "http://$ip:$ES_HTTP_PORT/_cat/health")
            if echo "$start_info" | grep -w 'green' &> /dev/null || echo "$start_info" | grep -w 'yellow' &> /dev/null; then
                ((healthy_node_count++))
                log_success "节点 $ip 已启动并健康"
            elif echo "$start_info" | grep 'missing authentication credentials' &> /dev/null; then
                ((healthy_node_count++))
                log_success "节点 $ip 已启动并健康（缺少认证凭证）"
            else
                log_warn "节点 $ip 尚未启动或不健康"
            fi
        done

        # 检查所有节点是否都已启动
        if [ "$healthy_node_count" -eq "$node_count" ]; then
            ES_HEALTH_STATUS="green"
            log_success "所有节点已启动并健康!"
            break
        fi

        sleep 20

    done
}

##############################################
# 初始化集群
##############################################
init_cluster() {

    # 仅在第一个服务器上生成证书，并将证书拷贝至另外两个服务器
    if [ "${ES_IP_LIST[0]}" = "$HOST_IP" ]; then
        log_info "将在第一台服务器 ${ES_IP_LIST[0]} 生成证书elastic-certificates.p12"
        "$ES_HOME"/bin/elasticsearch-certutil cert --days 36500 -out "$ES_HOME"/config/elastic-certificates.p12 -pass "" > /dev/null 2>&1
        chown "$ES_USER":"$ES_USER" -R "$ES_HOME"/config/elastic-certificates.p12
        for (( i = 1; i < ${#ES_IP_LIST[@]}; i++ )); do
            log_info "拷贝 elastic-certificates.p12 至 ${ES_IP_LIST[i]}:$DIR 目录"
            scp "$ES_HOME"/config/elastic-certificates.p12 root@"${ES_IP_LIST[i]}":"$DIR"
        done
    fi

    # 其余服务器等待证书拷贝
    if [ "${ES_IP_LIST[0]}" != "$HOST_IP" ]; then
        while :
        do  
            log_info "等待elastic-certificates.p12拷贝至当前服务器..."
            if [ -f "$DIR"/elastic-certificates.p12 ] && [ ! -f "$ES_HOME"/config/elastic-certificates.p12 ]; then
                log_info "移动 elastic-certificates.p12 至 ${ES_HOME}/config 目录"
                mv "$DIR"/elastic-certificates.p12 "$ES_HOME"/config &&
                chown "$ES_USER":"$ES_USER" -R "$ES_HOME"/config/elastic-certificates.p12
                break
            fi
            sleep 10
        done
    fi

    log_info "启动Elasticsearch"
    su - $ES_USER -c "$ES_HOME/bin/elasticsearch -d" &> /dev/null

    log_info "等待Elasticsearch所有节点启动"
    wait_es_start
    
    log_info "将在第一台服务器 ${ES_IP_LIST[0]} 设置集群密码"
    if [ "${ES_IP_LIST[0]}" = "$HOST_IP" ]; then
        if [ "$ES_HEALTH_STATUS" = "green" ]; then	
            while :
            do
                # auto 是自动生成密码 interactive为手动创建密码
                read -p "集群已经启动，请输入 Y 或 N 来设置密码，Y为自动生成密码，N为手动设置密码：" -r option
                case "$option" in
                    Y|y)
                        echo "自动生成密码，请保存密码"
                        "$ES_HOME"/bin/elasticsearch-setup-passwords auto
                        break
                    ;;
                    N|n)
                        echo "手动生成密码，请根据提示粘贴密码"
                        "$ES_HOME"/bin/elasticsearch-setup-passwords interactive
                        break
                    ;;
                    *)
                        echo "输入不合法：请输入Y或N"
                    ;;
                esac
            done
        fi

        while :
        do
            read -p "是否调整集群分片数? Y/N " -r item
            case "${item}" in
                Y|y)
                    log_info "将最大分片数调整为3000，请输入elastic账号对应的密码"
                    curl -u elastic -XPUT "http://127.0.0.1:9200/_cluster/settings" \
                        -H 'Content-Type: application/json' \
                        -d '{ "persistent" : { "cluster.max_shards_per_node" : "3000" } }'
                    log_info "调整集群分片数成功"
                    break
                ;;
                N|n)
                    log_info "不调整集群分片数"
                    break
                ;;
                *)
                    log_info "输入不合法：请输入Y或N"
                ;;
            esac
            
        done
    fi
}

main() {
    log_info "开始安装Elasticsearch集群"
    init_env
    install_elasticsearch
    add_service_config
    init_cluster
    log_success "Elasticsearch集群安装完成"
}

main